package kafka

import (
	"context"
	"github.com/Shopify/sarama"
	"time"
)

type (
	ConsumerMessageHandleFunc func(msg *sarama.ConsumerMessage) (err error) //消息处理函数 ,返回错误，消息不会提交
	Consumer                  struct {
		HandleFunc ConsumerMessageHandleFunc
	}

	//消费服务
	ConsumeService struct {
		Config *sarama.Config
		Addr   []string
	}
)

func (this *Consumer) Setup(s sarama.ConsumerGroupSession) error {
	return nil
}

func (this *Consumer) Cleanup(s sarama.ConsumerGroupSession) error {
	return nil
}

func (this *Consumer) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {

	for message := range claim.Messages() {

		for {
			if err := this.HandleFunc(message); err == nil {
				break
			}

			time.Sleep(3 * time.Second)
		}

		session.MarkMessage(message, "")
	}

	return nil
}

func NewConsumeService(addr []string) *ConsumeService {

	this := &ConsumeService{
		Addr: addr,
	}

	this.Config = sarama.NewConfig()
	this.Config.Version = sarama.V0_10_2_0
	this.Config.Consumer.Offsets.Initial = sarama.OffsetOldest
	this.Config.Consumer.Return.Errors = true

	//this.Config.Net.SASL.User = DefaultKafkaConfig.User
	//this.Config.Net.SASL.Password = DefaultKafkaConfig.Password
	//if this.Config.Net.SASL.User != "" {
	//	this.Config.Net.SASL.Enable = true
	//}

	return this
}

//消费信息
func (this *ConsumeService) Consume(topic string, consumerNum int, handleFun ConsumerMessageHandleFunc) {
	for i := 0; i < consumerNum; i++ {
		go this.newConsumerWithReconnection(topic, handleFun)
	}
}

//消费者组失联后重连
func (this *ConsumeService) newConsumerWithReconnection(topic string, handleFun ConsumerMessageHandleFunc) {
	var err error
	consumer := &Consumer{HandleFunc: handleFun}

	for {
		if err = this.newConsumer(topic, topic, consumer); err == nil {
			break
		}
		time.Sleep(5 * time.Second)
	}
}

//生成一个新的消费者
func (this *ConsumeService) newConsumer(topic string, groupId string, consumerHandler sarama.ConsumerGroupHandler) (err error) {

	var (
		consumerGroup sarama.ConsumerGroup
		client        sarama.Client
	)

	if client, err = sarama.NewClient(this.Addr, this.Config); err != nil {
		return
	}

	if consumerGroup, err = sarama.NewConsumerGroupFromClient(groupId, client); err != nil {
		return
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	if err := consumerGroup.Consume(ctx, []string{topic}, consumerHandler); err != nil {
		return err
	}

	defer consumerGroup.Close()

	for {
		select {
		case err = <-consumerGroup.Errors():
			return
		default:
			time.Sleep(1 * time.Second)
		}
	}

}

//获取topic 下分区数量
func (this *ConsumeService) getPartitionNum(topic string) (num int, err error) {

	var (
		client     sarama.Client
		consumer   sarama.Consumer
		partitions []int32
	)

	if client, err = sarama.NewClient(this.Addr, this.Config); err != nil {
		return
	}

	if consumer, err = sarama.NewConsumerFromClient(client); err != nil {
		return
	}

	if partitions, err = consumer.Partitions(topic); err != nil {
		return
	}

	return len(partitions), err
}
